#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd

# Dew Point Temperature calculation
def Tdcalc(H,p,T):
	ep = 0.622  #g/g
	e0 = 611.3  #Pa
	b = 17.2694 
	T1 = 273.15
	T2 = 35.86  
	Td = (b*T1-T2*xr.ufuncs.log(H*p/(ep*e0)))/(b-xr.ufuncs.log(H*p/(ep*e0)))-273.15
	Td = xr.where(T<Td,T,Td)
	return Td

# K-index calculation
def Kicalc(T850,T500,Td850,T700,Td700):
	Ki = (T850-T500)+Td850-(T700-Td700)
	return Ki	

year=int(sys.argv[1])

path_hus850='./ERA5/H850/'      # 850 hPa specific humidity in g/g
path_hus700='./ERA5/H700/'      # 700 hPa specific humidity in g/g
path_ta850='./ERA5/T850/'       # 850 hPa temperature in Kelvin
path_ta700='./ERA5/T700/'       # 700 hPa temperature in Kelvin
path_ta500='./ERA5/T500/'       # 500 hPa temperature in Kelvin
path_out_Ki='./K-index/'  

## 850 hPa specific humidity
hus850 = xr.open_dataset(path_hus850+'ERA5_H850_EU_'+str(year)+'.nc',decode_times=False)['q']
## 700 hPa specific humidity
hus700 = xr.open_dataset(path_hus700+'ERA5_H700_EU_'+str(year)+'.nc',decode_times=False)['q']
## 850 hPa air temperature
ta850 = xr.open_dataset(path_ta850+'ERA5_T850_EU_'+str(year)+'.nc',decode_times=False)['t']-273.15
## 700 hPa air temperature
ta700 = xr.open_dataset(path_ta700+'ERA5_T700_EU_'+str(year)+'.nc',decode_times=False)['t']-273.15
## 500 hPa air temperature
ta500 = xr.open_dataset(path_ta500+'ERA5_T500_EU_'+str(year)+'.nc',decode_times=False)['t']-273.15

# Creating time vector
dates = pd.date_range(start=str(year)+'-01-01', end=str(year)+'-12-31', freq='1D')
hus850['time'] = dates
hus700['time'] = dates
ta850['time'] = dates
ta700['time'] = dates
ta500['time'] = dates

# Loop for computing daily K-index
Ki = xr.full_like(ta850,np.nan)
for t in range(ta850.time.size):
	print(dates[t])

	H850 = hus850.isel(time=t).copy()
	H700 = hus700.isel(time=t).copy()
	T850 = ta850.isel(time=t).copy()
	T700 = ta700.isel(time=t).copy()
	T500 = ta500.isel(time=t).copy()

	Td850 = Tdcalc(H850,85000.,T850)
	Td700 = Tdcalc(H700,70000.,T700)
	Ki.values[t,:] = Kicalc(T850,T500,Td850,T700,Td700)

print('Year: {} \n Loop end. Saving files...'.format(str(year)))

Ki.name = 'Ki'
Ki.attrs['standard_name'] = 'Ki'
Ki.attrs['long_name'] = 'K-index'
Ki.attrs['units'] = 'Celsius'

# Saving files
Ki.to_netcdf(path_out_Ki+'Ki_'+str(year)+'.nc')
